package KafkaApps;

import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.LongDeserializer;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Properties;


public class ConsumerPiece {
    private final static String TOPIC = "testTrans";
    static String fileName = "trans_file/consumer_file/after_" + "piece" + ".txt";

    public static Consumer<Long, byte[]> createConsumer() {
        final Properties sets = new Properties();
        sets.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "master:9092");
        sets.put(ConsumerConfig.GROUP_ID_CONFIG, "transFile");
        sets.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,LongDeserializer.class.getName());
        sets.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());

        final Consumer<Long, byte[]> consumer = new KafkaConsumer<>(sets);
        consumer.subscribe((Collections.singletonList(TOPIC)));
        return consumer;
    }

    //未用到
    private static boolean writeFile(byte[] bytes){
        FileOutputStream fileOutputStream = null;
        try {
            fileOutputStream = new FileOutputStream(fileName);
            fileOutputStream.write(bytes);
            return true;
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                fileOutputStream.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return false;
    }

    private static void getMsg(Consumer<Long, byte[]> consumer){
        final int giveUp = 10;   int noRecordsCount = 0;
        FileOutputStream fileOutputStream = null;
        HashMap<Integer,byte[]> file = new HashMap<>();
        while (true) {
            final ConsumerRecords<Long, byte[]> consumerRecords =
                    consumer.poll(1000);

            if (consumerRecords.count()==0) {
                noRecordsCount++;
                if (noRecordsCount > giveUp) break;
                else continue;
            }

            for (ConsumerRecord<Long, byte[]> consumerRecord: consumerRecords) {
                for(Header kafkaHeader:consumerRecord.headers()){
                    if((kafkaHeader).key().equals("number")){
                        System.out.println(Integer.parseInt(new String(kafkaHeader.value())));
                        file.put(Integer.parseInt(new String(kafkaHeader.value())),consumerRecord.value());
                        break;
                    }
                }

            }



            consumer.commitAsync();
        }
        int count = file.size();
        try {
            fileOutputStream = new FileOutputStream(fileName);
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        }
        for (int i = 1; i <= count; i++){
            for (Iterator<Integer> it = file.keySet().iterator(); it.hasNext(); ) {
                int key = it.next();
                if (key == i){
                    //writeFile(file.get(key));
                    try {
                        System.out.println("正在写入"+i+"个");
                        fileOutputStream.write(file.get(key));
                        System.out.println("写入"+i+"个成功");
                        System.out.println();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            }
        }
        try {
            fileOutputStream.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        consumer.close();
        System.out.println("DONE");
    }

    public static void main(String... args)  {
        final Consumer<Long, byte[]> consumer = createConsumer();
        getMsg(consumer);
    }
}
